#!/usr/bin/env python3
"""
LangChain 智能体工具调用框架
功能：
1. 定义4个示例函数作为工具
2. 创建带记忆的智能体
3. 支持自然语言交互和工具选择
4. 使用现代化的记忆管理
"""

import os
from typing import Dict, Any
from dotenv import load_dotenv

# LangChain 核心模块
from langchain_core.tools import tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.messages import HumanMessage

# 语言模型 - 使用通义千问
from langchain_community.llms import Tongyi

# 智能体相关
from langchain.agents import create_react_agent, AgentExecutor

# 加载环境变量
load_dotenv()

# ==================== 工具函数定义 ====================

@tool
def func1(query: str) -> str:
    """处理数学计算相关的问题"""
    return f"数学计算工具处理: {query}\n结果: 这是一个数学计算的示例结果"

@tool
def func2(query: str) -> str:
    """处理天气查询相关的问题"""
    return f"天气查询工具处理: {query}\n结果: 今天天气晴朗，温度25°C"

@tool
def func3(query: str) -> str:
    """处理翻译相关的问题"""
    return f"翻译工具处理: {query}\n结果: 这是翻译后的内容示例"

@tool
def func4(query: str) -> str:
    """处理搜索相关的问题"""
    return f"搜索工具处理: {query}\n结果: 找到了相关的搜索结果"

# ==================== 智能体类 ====================

class LangChainAgent:
    """LangChain 智能体框架类"""
    
    def __init__(self):
        self.llm = None
        self.tools = [func1, func2, func3, func4]
        self.agent_executor = None
        self.store = {}  # 会话存储
        
    def setup_llm(self):
        """初始化语言模型"""
        print("正在初始化语言模型...")
        
        api_key = os.getenv("DASHSCOPE_API_KEY")
        if not api_key:
            raise ValueError("请在.env文件中设置DASHSCOPE_API_KEY")
            
        self.llm = Tongyi(
            model="qwen-turbo",
            api_key=api_key,
            temperature=0.7,
            max_tokens=1024
        )
        print("语言模型初始化完成")
        
    def setup_agent(self):
        """创建智能体"""
        print("正在创建智能体...")
        
        # 创建标准的ReAct提示模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个智能助手，可以使用工具来帮助用户。"),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", """
你有以下工具可以使用:
{tools}

使用以下格式:

Question: 用户的问题
Thought: 你应该思考要做什么
Action: 要采取的行动，应该是 [{tool_names}] 中的一个
Action Input: 行动的输入
Observation: 行动的结果
... (这个 Thought/Action/Action Input/Observation 可以重复N次)
Thought: 我现在知道最终答案了
Final Answer: 对原始问题的最终答案

开始!

Question: {input}
Thought:{agent_scratchpad}""")
        ])
        
        # 创建基础智能体
        agent = create_react_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=prompt
        )
        
        # 创建智能体执行器
        base_executor = AgentExecutor(
            agent=agent,
            tools=self.tools,
            verbose=True,
            handle_parsing_errors=True,
            max_iterations=3
        )
        
        # 添加记忆功能
        self.agent_executor = RunnableWithMessageHistory(
            base_executor,
            self.get_session_history,
            input_messages_key="input",
            history_messages_key="chat_history",
        )
        
        print("智能体创建完成")
        
    def get_session_history(self, session_id: str) -> InMemoryChatMessageHistory:
        """获取会话历史"""
        if session_id not in self.store:
            self.store[session_id] = InMemoryChatMessageHistory()
        return self.store[session_id]
        
    def initialize(self):
        """初始化智能体"""
        print("=== 初始化 LangChain 智能体框架 ===")
        self.setup_llm()
        self.setup_agent()
        print("\n=== 初始化完成 ===")
        print(f"可用工具: {[tool.name for tool in self.tools]}")
        
    def chat(self, user_input: str, session_id: str = "default") -> str:
        """与智能体对话"""
        try:
            response = self.agent_executor.invoke(
                {"input": user_input},
                config={"configurable": {"session_id": session_id}}
            )
            return response["output"]
        except Exception as e:
            return f"处理请求时出错: {str(e)}"
            
    def interactive_chat(self):
        """交互式对话"""
        print("\n=== 开始交互式对话 ===")
        print("输入 'quit' 或 'exit' 退出对话")
        print("输入 'clear' 清除对话历史")
        print("-" * 50)
        
        session_id = "user_session"
        
        while True:
            try:
                user_input = input("\n用户: ").strip()
                
                if user_input.lower() in ['quit', 'exit', '退出']:
                    print("再见！")
                    break
                    
                if user_input.lower() in ['clear', '清除']:
                    if session_id in self.store:
                        del self.store[session_id]
                    print("对话历史已清除")
                    continue
                    
                if not user_input:
                    continue
                    
                print("\n助手: ", end="")
                response = self.chat(user_input, session_id)
                print(response)
                
            except KeyboardInterrupt:
                print("\n\n对话已中断")
                break
            except Exception as e:
                print(f"\n发生错误: {e}")

# ==================== 测试函数 ====================

def test_agent():
    """测试智能体功能"""
    print("=== 开始测试 ===")
    
    # 创建智能体实例
    agent = LangChainAgent()
    agent.initialize()
    
    # 测试对话
    test_cases = [
        "你好，我是小明",
        "帮我计算 2+3 等于多少",
        "今天天气怎么样？",
        "请翻译 'Hello World' 为中文",
        "搜索一下人工智能的相关信息",
        "你还记得我的名字吗？"
    ]
    
    session_id = "test_session"
    
    for i, test_input in enumerate(test_cases, 1):
        print(f"\n--- 测试 {i} ---")
        print(f"用户: {test_input}")
        response = agent.chat(test_input, session_id)
        print(f"助手: {response}")
        print("-" * 30)
    
    return agent

# ==================== 主函数 ====================

def main():
    """主函数"""
    try:
        # 运行测试
        agent = test_agent()
        
        # 启动交互式对话
        agent.interactive_chat()
        
    except Exception as e:
        print(f"程序运行出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()